import struct,numpy as np

# image_filename，idx格式的图片文件名"train-images-idx3-ubyte"，共60000张
# label_filename，标签文件名"train-labels-idx1-ubyte"，与图片文件对应，亦为60000张
# batch_size，生成器每次提供的数据量，即每批次多少图片+标签进行训练，一般为32
# drop_last，如果样本数不够batch_size整除，最后一组不够批量的样本是否丢弃？

def dataReader(image_filename, label_filename, batch_size, drop_last=False):
    with open(image_filename,'rb') as f:
        buff = f.read()
        # 解析文件头信息，依次为魔数、图片数量、每张图片高、每张图片宽
        offset = 0
        fmt = '>iiii'
        magic_number, num_images, num_rows, num_cols = struct.unpack_from(fmt, buff, offset)
        print(image_filename+'魔数:%d, 图片数量: %d张, 图片大小: %d*%d' % (magic_number, num_images, num_rows, num_cols))

        # 解析数据集
        image_size = num_rows * num_cols
        offset += struct.calcsize(fmt)
        fmt_image = '>' + str(image_size) + 'B'
        images = np.empty((num_images, num_rows, num_cols))
        for i in range(num_images):
            # if (i + 1) % 10000 == 0:
            #     print('已解析 %d' % (i + 1) + '张')
            # yield np.array(struct.unpack_from(fmt_image, buff, offset)).reshape((num_rows, num_cols))
            images[i] = np.array(struct.unpack_from(fmt_image, buff, offset)).reshape((num_rows, num_cols))
            offset += struct.calcsize(fmt_image)


    with open(label_filename, 'rb') as f2:
        buff2 = f2.read()
        # 解析文件头信息，依次为魔数、图片数量、每张图片高、每张图片宽
        offset2 = 0
        fmt2 = '>ii'
        magic_number2, num_images2 = struct.unpack_from(fmt2, buff2, offset2)
        print(label_filename+'魔数:%d, 图片数量: %d张' % (magic_number2, num_images2))

        # 解析数据集
        offset2 += struct.calcsize(fmt2)
        fmt_image2 = '>B'
        lables = np.empty(num_images2)
        for i in range(num_images2):
            # if (i + 1) % 10000 == 0:
            #     print('已解析 %d' % (i + 1) + '张')
            # yield np.array(struct.unpack_from(fmt_image, buff, offset)).reshape((num_rows, num_cols))
            lables[i] = struct.unpack_from(fmt_image2, buff2, offset2)[0]
            offset2 += struct.calcsize(fmt_image2)

    return images,lables

# dataReader('train-images.idx3-ubyte',
#         'train-labels-idx1-ubyte',
#         32,True)

def train():
    for batch_id, data in enumerate(
            dataReader('train-images.idx3-ubyte','train-labels.idx1-ubyte',32,True)):
        print(data)

train()

def test():
    pass

def tarin_test(batch_size):
    n = 0
    while n < (60000%batch_size):
        yield test()